// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#pragma once

#include <cstdint>
#include <memory>
#include <string>
#include <vector>

#include <nebula/array/array_nested.h>
#include <nebula/array/builder_base.h>
#include <nebula/array/data.h>
#include <nebula/core/buffer_builder.h>
#include <nebula/core/memory_pool.h>
#include <turbo/utility/status.h>
#include <nebula/types/type.h>
#include <turbo/base/macros.h>

namespace nebula {

    /// \addtogroup nested-builders
    ///
    /// @{

    /// \brief Base class for union array builds.
    ///
    /// Note that while we subclass ArrayBuilder, as union types do not have a
    /// validity bitmap, the bitmap builder member of ArrayBuilder is not used.
    class TURBO_EXPORT BasicUnionBuilder : public ArrayBuilder {
    public:
        turbo::Status finish_internal(std::shared_ptr<ArrayData> *out) override;

        /// \cond FALSE
        using ArrayBuilder::finish;
        /// \endcond

        turbo::Status finish(std::shared_ptr<UnionArray> *out) { return finish_typed(out); }

        /// \brief Make a new child builder available to the UnionArray
        ///
        /// \param[in] new_child the child builder
        /// \param[in] field_name the name of the field in the union array type
        /// if type inference is used
        /// \return child index, which is the "type" argument that needs
        /// to be passed to the "append" method to add a new element to
        /// the union array.
        int8_t AppendChild(const std::shared_ptr<ArrayBuilder> &new_child,
                           const std::string &field_name = "");

        std::shared_ptr<DataType> type() const override;

        int64_t length() const override { return types_builder_.length(); }

    protected:
        BasicUnionBuilder(MemoryPool *pool, int64_t alignment,
                          const std::vector<std::shared_ptr<ArrayBuilder>> &children,
                          const std::shared_ptr<DataType> &type);

        int8_t NextTypeId();

        std::vector<std::shared_ptr<Field>> child_fields_;
        std::vector<int8_t> type_codes_;
        UnionMode::type mode_;

        std::vector<ArrayBuilder *> type_id_to_children_;
        std::vector<int> type_id_to_child_id_;
        // for all type_id < dense_type_id_, type_id_to_children_[type_id] != nullptr
        int8_t dense_type_id_ = 0;
        TypedBufferBuilder<int8_t> types_builder_;
    };

    /// \class DenseUnionBuilder
    ///
    /// This API is EXPERIMENTAL.
    class TURBO_EXPORT DenseUnionBuilder : public BasicUnionBuilder {
    public:
        /// Use this constructor to initialize the UnionBuilder with no child builders,
        /// allowing type to be inferred. You will need to call AppendChild for each of the
        /// children builders you want to use.
        explicit DenseUnionBuilder(MemoryPool *pool,
                                   int64_t alignment = kDefaultBufferAlignment)
                : BasicUnionBuilder(pool, alignment, {}, dense_union(FieldVector{})),
                  offsets_builder_(pool, alignment) {}

        /// Use this constructor to specify the type explicitly.
        /// You can still add child builders to the union after using this constructor
        DenseUnionBuilder(MemoryPool *pool,
                          const std::vector<std::shared_ptr<ArrayBuilder>> &children,
                          const std::shared_ptr<DataType> &type,
                          int64_t alignment = kDefaultBufferAlignment)
                : BasicUnionBuilder(pool, alignment, children, type),
                  offsets_builder_(pool, alignment) {}

        turbo::Status append_null() final {
            const int8_t first_child_code = type_codes_[0];
            ArrayBuilder *child_builder = type_id_to_children_[first_child_code];
            TURBO_RETURN_NOT_OK(types_builder_.append(first_child_code));
            TURBO_RETURN_NOT_OK(
                    offsets_builder_.append(static_cast<int32_t>(child_builder->length())));
            // append a null arbitrarily to the first child
            return child_builder->append_null();
        }

        turbo::Status append_nulls(int64_t length) final {
            const int8_t first_child_code = type_codes_[0];
            ArrayBuilder *child_builder = type_id_to_children_[first_child_code];
            TURBO_RETURN_NOT_OK(types_builder_.append(length, first_child_code));
            TURBO_RETURN_NOT_OK(
                    offsets_builder_.append(length, static_cast<int32_t>(child_builder->length())));
            // append just a single null to the first child
            return child_builder->append_null();
        }

        turbo::Status append_empty_value() final {
            const int8_t first_child_code = type_codes_[0];
            ArrayBuilder *child_builder = type_id_to_children_[first_child_code];
            TURBO_RETURN_NOT_OK(types_builder_.append(first_child_code));
            TURBO_RETURN_NOT_OK(
                    offsets_builder_.append(static_cast<int32_t>(child_builder->length())));
            // append an empty value arbitrarily to the first child
            return child_builder->append_empty_value();
        }

        turbo::Status append_empty_values(int64_t length) final {
            const int8_t first_child_code = type_codes_[0];
            ArrayBuilder *child_builder = type_id_to_children_[first_child_code];
            TURBO_RETURN_NOT_OK(types_builder_.append(length, first_child_code));
            TURBO_RETURN_NOT_OK(
                    offsets_builder_.append(length, static_cast<int32_t>(child_builder->length())));
            // append just a single empty value to the first child
            return child_builder->append_empty_value();
        }

        /// \brief append an element to the UnionArray. This must be followed
        ///        by an append to the appropriate child builder.
        ///
        /// \param[in] next_type type_id of the child to which the next value will be appended.
        ///
        /// The corresponding child builder must be appended to independently after this method
        /// is called.
        turbo::Status append(int8_t next_type) {
            TURBO_RETURN_NOT_OK(types_builder_.append(next_type));
            if (type_id_to_children_[next_type]->length() == kListMaximumElements) {
                return turbo::resource_exhausted_error(
                        "a dense UnionArray cannot contain more than 2^31 - 1 elements from a single "
                        "child");
            }
            auto offset = static_cast<int32_t>(type_id_to_children_[next_type]->length());
            return offsets_builder_.append(offset);
        }

        turbo::Status append_array_slice(const ArraySpan &array, int64_t offset,
                                       int64_t length) override;

        turbo::Status finish_internal(std::shared_ptr<ArrayData> *out) override;

    private:
        TypedBufferBuilder<int32_t> offsets_builder_;
    };

    /// \class SparseUnionBuilder
    ///
    /// This API is EXPERIMENTAL.
    class TURBO_EXPORT SparseUnionBuilder : public BasicUnionBuilder {
    public:
        /// Use this constructor to initialize the UnionBuilder with no child builders,
        /// allowing type to be inferred. You will need to call AppendChild for each of the
        /// children builders you want to use.
        explicit SparseUnionBuilder(MemoryPool *pool,
                                    int64_t alignment = kDefaultBufferAlignment)
                : BasicUnionBuilder(pool, alignment, {}, sparse_union(FieldVector{})) {}

        /// Use this constructor to specify the type explicitly.
        /// You can still add child builders to the union after using this constructor
        SparseUnionBuilder(MemoryPool *pool,
                           const std::vector<std::shared_ptr<ArrayBuilder>> &children,
                           const std::shared_ptr<DataType> &type,
                           int64_t alignment = kDefaultBufferAlignment)
                : BasicUnionBuilder(pool, alignment, children, type) {}

        /// \brief append a null value.
        ///
        /// A null is appended to the first child, empty values to the other children.
        turbo::Status append_null() final {
            const auto first_child_code = type_codes_[0];
            TURBO_RETURN_NOT_OK(types_builder_.append(first_child_code));
            TURBO_RETURN_NOT_OK(type_id_to_children_[first_child_code]->append_null());
            for (int i = 1; i < static_cast<int>(type_codes_.size()); ++i) {
                TURBO_RETURN_NOT_OK(type_id_to_children_[type_codes_[i]]->append_empty_value());
            }
            return turbo::OkStatus();
        }

        /// \brief append multiple null values.
        ///
        /// Nulls are appended to the first child, empty values to the other children.
        turbo::Status append_nulls(int64_t length) final {
            const auto first_child_code = type_codes_[0];
            TURBO_RETURN_NOT_OK(types_builder_.append(length, first_child_code));
            TURBO_RETURN_NOT_OK(type_id_to_children_[first_child_code]->append_nulls(length));
            for (int i = 1; i < static_cast<int>(type_codes_.size()); ++i) {
                TURBO_RETURN_NOT_OK(
                        type_id_to_children_[type_codes_[i]]->append_empty_values(length));
            }
            return turbo::OkStatus();
        }

        turbo::Status append_empty_value() final {
            TURBO_RETURN_NOT_OK(types_builder_.append(type_codes_[0]));
            for (int8_t code: type_codes_) {
                TURBO_RETURN_NOT_OK(type_id_to_children_[code]->append_empty_value());
            }
            return turbo::OkStatus();
        }

        turbo::Status append_empty_values(int64_t length) final {
            TURBO_RETURN_NOT_OK(types_builder_.append(length, type_codes_[0]));
            for (int8_t code: type_codes_) {
                TURBO_RETURN_NOT_OK(type_id_to_children_[code]->append_empty_values(length));
            }
            return turbo::OkStatus();
        }

        /// \brief append an element to the UnionArray. This must be followed
        ///        by an append to the appropriate child builder.
        ///
        /// \param[in] next_type type_id of the child to which the next value will be appended.
        ///
        /// The corresponding child builder must be appended to independently after this method
        /// is called, and all other child builders must have null or empty value appended.
        turbo::Status append(int8_t next_type) { return types_builder_.append(next_type); }

        turbo::Status append_array_slice(const ArraySpan &array, int64_t offset,
                                       int64_t length) override;
    };

    /// @}

}  // namespace nebula
